#!/usr/bin/python
from pylab import *

# import numpy as np
# import matplotlib.pyplot as plt
# import matplotlib
from matplotlib.collections import PatchCollection
# import matplotlib.path as mpath
import matplotlib.patches as mpatches
import matplotlib.lines as mlines
import copy

# font = "sans-serif"
fig = figure(figsize=(70,10))
# ax = fig.add_axes([0,0,1,1])  # the ticks not displayed
ax = fig.add_axes([0.1,0.1,0.8,0.8])
# ax = fig.add_subplot(111)

# fig, ax = plt.subplots()


# x = arange(0,1,0.01)
# plot(x,sin(2*pi*x))


# # create 3x3 grid to plot the artists
# # pos = np.mgrid[0.2:0.8:3j, 0.2:0.8:3j].reshape(2, -1)

# patches = []

# # add a circle
# # art = mpatches.Circle(pos[:,0], 0.1,ec="none")
# art = mpatches.Circle((0.5,0.5), 0.1,ec="none")
# patches.append(art)
# # ax.text(0.5, 0.5-0.15, "Circle", ha="center", family=font, size=14)

# art = mpatches.Rectangle((0.5,0.5), 0.5, 0.2, ec='none')
# # art = mpatches.Rectangle((0.5,0.5), 0.5, 0.2, fc='r', ec='none')
# patches.append(art)

# art = mpatches.Arrow(0,0,0.3,0.3,width=0.2)
# patches.append(art)

# # lines (not append to patches, added to ax)
# line = mlines.Line2D([0,0.2],[0.3,0.4], lw=4, color='r', alpha=0.4) # draw a line (0,0.3) - (0.2,0.4)
# lineX = linspace(0,2,5)
# lineY = [0,0,1,0,0]
# line = mlines.Line2D(lineX, lineY, lw=4, c='b', alpha=0.4)

# setp(line, color='g')           # set the properties of 'line'


# # colors = 100*np.random.rand(len(patches))
# collection = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=0.4)
# # collection.set_array(np.array(colors))
# # ax.add_collection(collection)


# ax.add_line(line)               # add lines
xlim(0,2100)                       # must set the xlim/ylim manually
ylim(0,300)
# ax.set_xticks([])
# ax.set_yticks([])
ax.set_xticks([-1, -0.5, 0, 0.5, 1], ['$-\pi$', '$-\pi/2$', '$0$', '$+\pi/2$', '$+\pi$'])


#def plot_circle(center, radius, fill=True, color="black", **kwargs):
# def plot_circle(center, radius, **kwargs):
#     patch = mpatches.Circle(center, radius, kwargs=args)
#     ax.add_patch(patch)

 # def plot_circle(center, radius, **kwargs):
 #    patch = mpatches.Circle(center, radius, kwargs=args)
 #    ax.add_patch(patch)
   
def plot_circle(center, radius, **kwargs):
   patch = mpatches.Circle(center, radius, **kwargs)
   ax.add_patch(patch)

def plot_half_circle(center, radius, theta1, theta2, **kwargs):
   patch = mpatches.Wedge(center, radius, theta1, theta2, **kwargs)
   ax.add_patch(patch)
   
def plot_dot(pos, **kwargs):
   radius = 2
   linewidth = 1.5
   color = "black"
   fill = "True"

   if "color" in kwargs:
      color = kwargs['color']
   if "fill" in kwargs:
      fill = kwargs['fill']

   # plot_circle(pos, radius, fill=True, color=dotcolor, alpha=0.5)
   # plot_circle(pos, radius, fill=False, color="red", alpha=0.5)
   dotstyle = "solidcircle"
   if "dotstyle" in kwargs:
      dotstyle = kwargs["style"]

   # if  style == "solidcircle":
   #    plot_circle(pos, radius, color=color, fill=fill)
   if dotstyle == "halfcircle":
      plot_half_circle(pos, radius, theta1=90, theta2=270, color=color, fill=fill)
   else:
      plot_circle(pos, radius, color=color, fill=fill, linewidth=linewidth)
   
      

   
   


       
def plot_rect(start_pos, end_pos, **kwargs):
   # start_pos and end_pos any pair of the diagonal point
   if start_pos[0] < end_pos[0]:
      sx = start_pos[0]
      ex = end_pos[0]
   else:
      sx = end_pos[0]
      ex = start_pos[0]

   if start_pos[1] < end_pos[1]:
      sy = start_pos[1]
      ey = end_pos[1]
   else:
      sy = end_pos[1]
      ey = start_pos[1]

   width = ex - sx + 1
   height = ey - sy + 1
   
   # sx = start_pos[0] if start_pos[0]<end_pos[0] else end_pos[0]
   # sy = start_pos[1] if start_pos[1]<end_pos[1] else end_pos[1]
   
   patch = mpatches.Rectangle((sx, sy), width, height, **kwargs)
   ax.add_patch(patch)


# ==================== parameters and global vars ====================
# width = 8
width = 16                      # 16x16 multiply
# width = 32

col_middle = width
pps_num = width/2+1                      # num of pps
pps_width = width + 1                    # the width of last PP is width
cols_num = width*2
rows_num = width/2+1

# dots[rows_num][cols_num]
# dots = [[None]*cols_num for i in range(rows_num)]


step_x = 10                      # the diameter of a node in the grid
step_y = 10
# pos_base = (20,20)

# # def display_dots(pos_base=(0,0), step_x=10, step_y=10):
# def display_dots(dots, pos_base=(0,0), dotcolor="black"):
#    point_base = (pos_base[0]+step_x/2, pos_base[1]+step_y/2)
#    # for y in range(rows_num):
#    #     for x in range(cols_num):
#    #         pos = (pos_base[0]+step_x*(cols_num-x), pos_base[1]+step_y*(rows_num-y))
#    #         if dots[y][x] == None:
#    #             continue
#    #         elif dots[y][x] == 1:
#    #             plot_circle(pos, radius)
#    #         else:
#    #             text(pos[0], pos[1], dots[y][x], ha="center", va="center", family="sans-serif", size=10)

#    # add one row for display ticks
#    disp_rows_num = rows_num + 1
#    ticks = [str(i) for i in range(cols_num)]
#    disp_dots = [ticks] + dots
#    for y in range(disp_rows_num):
#       for x in range(cols_num):
#          pos = (point_base[0]+step_x*(cols_num-1-x), point_base[1]+step_y*(disp_rows_num-1-y))
#          if disp_dots[y][x] == 0:
#             continue
#          elif disp_dots[y][x] == 1:
#             plot_dot(pos, dotcolor)
#          else:
#             text(pos[0], pos[1], disp_dots[y][x], ha="center", va="center", family="sans-serif", size=10)

#    # draw the grid
#    pos_left_buttom = (pos_base[0], pos_base[1])
#    pos_right_top = (pos_base[0]+step_x*cols_num, pos_base[1]+step_y*disp_rows_num)
#    # the boundaries
#    plot_rect(pos_left_buttom, pos_right_top, fill=False, color="black")
#    # the table header
#    tab_header_start_pos = (pos_left_buttom[0], pos_right_top[1]-step_y)
#    tab_header_end_pos = pos_right_top
#    plot_rect(tab_header_start_pos, tab_header_end_pos, fill=True, color='red', alpha=0.8)

#    # the columns
#    # lineX = [pos_base[0]+step_x*i for i in range(1,cols_num)]
#    for i in range(1,cols_num):
#       x = pos_base[0]+step_x*i
#       y0 = pos_left_buttom[1]
#       y1 = pos_right_top[1]
#       line = mlines.Line2D([x,x], [y0, y1], lw=1, c='black')
#       ax.add_line(line)
#    # the rows
#    for i in range(1, disp_rows_num):
#       y = pos_base[1]+step_y*i
#       x0 = pos_left_buttom[0]
#       x1 = pos_right_top[0]
#       line = mlines.Line2D([x0,x1], [y,y], lw=1, c='black')
#       ax.add_line(line)
# 
      
# # the dots (PPs)
# for row in range(rows_num):
#     start_x = row*2
#     if row == rows_num-1:       # the last PP has one column less than others
#         end_x = start_x + pps_width-1
#     else:
#         end_x = start_x + pps_width
#     for col in range(start_x, end_x):
#         dots[row][col] = 1   # display the dot
# # the S[]
# for col in range(0, col_middle, 2):
#     dots[1+col/2][col] = "S[%d]" % (col/2)
# # the E[] & 1
# for col in range(col_middle+3, cols_num-1+2, 2):
#     row = (col-col_middle-3)/2+1
#     dots[row][col] = "E[%d]" % row
#     if col < cols_num-1:
#         dots[row][col+1] = "1"
# dots[0][col_middle+1] = "{~E[0]}"
# dots[0][col_middle+2] = "{~E[0]}"
# dots[0][col_middle+3] = "E[0]"

## ====================================================================================================
# def gen_empty_dots_array(rows, cols):
#    # return [[("disp","value")]*cols for i in range(rows)]
#    return [[0]*cols for i in range(rows)]

## ====================================================================================================



   
####################################################################################################
def gen_empty_dots_array(rows, cols):
   # return [[0]*cols for i in range(rows)]
   return [[(0,False)]*cols for i in range(rows)] # for each unit, two values : (val, isDispText), the value and display text or not (otherwise display dots)


# style : None:the original style(controlled by dots[x][x][1]); 
#         text:display all dots' value;
#         dot :display all dots as a dot (e.g, solid circle)
def display_dots(dots, pos_base=(0,0), style=None, **kwargs):
   point_base = (pos_base[0]+step_x/2, pos_base[1]+step_y/2)
   # add one row for display ticks
   disp_rows_num = rows_num + 1
   ticks = [(str(i),True) for i in range(cols_num)]
   disp_dots = [ticks] + dots
   for y in range(disp_rows_num):
      for x in range(cols_num):
         pos = (point_base[0]+step_x*(cols_num-1-x), point_base[1]+step_y*(disp_rows_num-1-y))
         assert(type(disp_dots[y][x]) == tuple)
         if disp_dots[y][x][0] == 0: # empty cell, do nothing, ignore the isDispText field
            continue
         elif (style=="dot" and y!=0) or (style!="text" and disp_dots[y][x][1]==False): # row 0 must display as text (it's ticks)
            plot_dot(pos, **kwargs)
         else:
            text(pos[0], pos[1], disp_dots[y][x][0], ha="center", va="center", family="sans-serif", size=10)
   # draw the grid
   pos_left_buttom = (pos_base[0], pos_base[1])
   pos_right_top = (pos_base[0]+step_x*cols_num, pos_base[1]+step_y*disp_rows_num)
   # the boundaries
   plot_rect(pos_left_buttom, pos_right_top, fill=False, color="black")
   # the table header
   tab_header_start_pos = (pos_left_buttom[0], pos_right_top[1]-step_y)
   tab_header_end_pos = pos_right_top
   plot_rect(tab_header_start_pos, tab_header_end_pos, fill=True, color='red', alpha=0.8)
   # the columns
   # lineX = [pos_base[0]+step_x*i for i in range(1,cols_num)]
   for i in range(1,cols_num):
      x = pos_base[0]+step_x*i
      y0 = pos_left_buttom[1]
      y1 = pos_right_top[1]
      line = mlines.Line2D([x,x], [y0, y1], lw=1, c='black')
      ax.add_line(line)
   # the rows
   for i in range(1, disp_rows_num):
      y = pos_base[1]+step_y*i
      x0 = pos_left_buttom[0]
      x1 = pos_right_top[0]
      line = mlines.Line2D([x0,x1], [y,y], lw=1, c='black')
      ax.add_line(line)


# def gen_pp(mcand, mplier_width, sign, prefix=""):
#    # assert(len(mplier)%2 == 0)   # the length of mplier must be an even num
#    # mplier_ext = [0] + mplier + [sign, sign] # 
#    # pps_num = len(mplier)/2+1
#    pps_num = mplier_width/2+1
#    pps_width = len(mcand) + 1

#    dots = gen_empty_dots_array(rows_num, cols_num)
   
#    for row in range(pps_num):
#       if row == pps_num - 1:    # the last row, one less than others
#          for i,v in enumerate(mcand):
#             if v == 1:
#                dots[row][row*2+i] = ("%spp[%d][%d]" % (prefix, row, i), False)
#       else:
#          for i,v in enumerate([mcand[-1]] + mcand):
#             if v == 1:
#                dots[row][row*2+i] = ("%spp[%d][%d]" % (prefix, row, i), False)
            
#       if row < pps_num-1:
#          dots[row+1][row*2] = ("%sS[%d]" % (prefix, row), True)
#       if row>=1 and row<=pps_num-2:
#          dots[row][row*2+pps_width] = ("%sE[%d]" % (prefix, row), True)
#       if row>=1 and row<=pps_num-3:
#          dots[row][row*2+pps_width+1] = ("1", True)
#    dots[0][pps_width] = ("~%sE[0]" % prefix, True) # for the first row
#    dots[0][pps_width+1] = ("~%sE[0]" % prefix, True)
#    dots[0][pps_width+2] = ("%sE[0]" % prefix, True)
#    return dots
   
# pp0_height_reduction should change only for counteracting or demo in few scheme
# pp_index_start only used for special scheme
def gen_pp(mcand_width, mplier_width, prefix="", pp_index_start=(0,0), pp0_height_reduction=True):
   assert(mplier_width%2 == 0)   # the length of mplier must be an even num

   pps_width = mcand_width + 1
   mplier_ext_width = mplier_width + 2 + 1
   pps_num = mplier_width/2+1
   
   dots = gen_empty_dots_array(rows_num, cols_num)
   
   for row in range(pps_num):
      if row == pps_num - 1:    # the last row, one less than others
         for i in range(pps_width-1):
            dots[row][row*2+i] = ("%spp[%d][%d]" % (prefix, pp_index_start[0]+row, pp_index_start[1]+i), False)
      else:
         for i in range(pps_width):
            dots[row][row*2+i] = ("%spp[%d][%d]" % (prefix, pp_index_start[0]+row, pp_index_start[1]+i), False)
            
      if row < pps_num-1:
         dots[row+1][row*2] = ("%sS[%d]" % (prefix, pp_index_start[0]+row), True)
      if row>=1 and row<=pps_num-2:
         dots[row][row*2+pps_width] = ("%sE[%d]" % (prefix, pp_index_start[0]+row), True)
      if row>=1 and row<=pps_num-3:
         dots[row][row*2+pps_width+1] = ("1'b1", True)
   dots[0][pps_width] = ("{~%sE[%d]}" % (prefix, pp_index_start[0]+0), True) # for the first row
   dots[0][pps_width+1] = ("{~%sE[%d]}" % (prefix,pp_index_start[0]+0), True)
   dots[0][pps_width+2] = ("%sE[%d]" % (prefix,pp_index_start[0]+0), True)

   # only for counteracting or demo in few scheme
   if pp0_height_reduction == False:
      assert(rows_num >= pps_num+1)     # must ensure rows_num is large enough (rows_num >= pps_num+1)
      for row in range(rows_num-1, 0, -1):
         for col in range(cols_num):
            dots[row][col] = dots[row-1][col];
      for col in range(cols_num):
         dots[0][col] = (0,False)
      dots[0][pps_width] = ("%sE[%d]" % (prefix, pp_index_start[0]+0), True)
      dots[1][pps_width] = ("1'b1", True)
      dots[1][pps_width+1] = ("1'b1", True)
      dots[1][pps_width+2] = (0, False)

   return dots


def set_dots_cell(dots, row, col, val):
   dots[row][col] = val

def set_dots_col(dots, col, col_val_list):
   for row, val in enumerate(col_val_list): 
      dots[row][col] = val

# notice : this func only set value but never clear any cells, so you may need input an empty dots as 'dots' param
# set pp[row] and corresponding S, if the length of input is less than the cols', then clear the rest of that row
def set_pp(dots, row, pp_val, s_val):
   start_col = row*2
   for col,val in enumerate(pp_val):
      dots[row][start_col+col] = val
   # for i in range(start_col):
   #    dots[row][i] = ('y', True)
   # for i in range(start_col+len(pp_val), cols_num):
   #    dots[row][i] = ('x', True)
   if s_val != None:
      dots[row+1][start_col] = s_val

# # set the display style : all dots show as text or a dot(e.g., solid circle)
# def set_dots_display_style(dots, style=None):
#    if style == "text":
#       text_option = True
#    elif style == "dot":
#       text_option = False
#    else:
#       return
      
#    for row in range(rows_num):
#       for col in range(cols_num):
#          if dots[row][col][0] != 0:
#             dots[row][col] = (dots[row][col][0], text_option)


def reorg_upward(dots):
   dots_up = gen_empty_dots_array(rows_num, cols_num)
   for col in range(cols_num):
      items = []
      for row in range(rows_num):
         if dots[row][col][0] != 0:
            items.append(dots[row][col])
      # items += [None]*(rows_num - len(items))
      for row, val in enumerate(items):
         dots_up[row][col] = val
   return dots_up

def get_col_items(dots):
   cols = []
   for col in range(cols_num):
      items = []
      for row in range(rows_num):
         val = dots[row][col][0]
         if val != 0:
            items.append(val)
      if items == []:
         items = ["0"]
      cols.append(items)
   return cols

def print_col_formula(dots):
   cols = get_col_items(dots)
   for col,items in enumerate(cols):
      print "col[%d] = %s;" % (col, ' + '.join(items))
            
      

def wallace_init(dots):
   global compress_array, compress_stage, compress_dots
   compress_array = []
   compress_stage = len(compress_array) - 1
   compress_dots = dots

def is_compress_finish():
   for col in range(cols_num):
      cnt = 0
      for row in range(rows_num):
         if compress_dots[row][col][0] != 0:
            cnt += 1
         else:
            break
      if cnt > 2:
         return False
   return True

# or the wallace tree compress split into two stages : select sets and compress
  
# the dots must already reorganized upwards as an inverted triangle
def wallace_compress():
   global compress_array, compress_stage, compress_dots

   compress_array.append([[] for i in range(cols_num)])
   compress_stage = len(compress_array) - 1

   cols_carry = [[] for i in range(cols_num)] # the carry of all columns
   compress_groups = []         # the groups to compress for all columns (the return value)
   
   for col in range(cols_num):
      items = []
      for row in range(rows_num):
         if compress_dots[row][col][0] == 0: # FIXME: but if there not exist any dots ?
            break               # finish
         else:
            items.append(compress_dots[row][col][0])
      n_compress_3_2 = len(items) / 3
      n_reminder = len(items) % 3
      
      curr_col_result = []
      fields = []
      # clac current col's sum
      i = -1
      for i in range(n_compress_3_2):
         comp_in = items[i*3:(i+1)*3]       # from i*3 to i*3+2
         comp_sum = "S%d_%d_%d" % (compress_stage+1, col, i) # stage_col_row
         # comp_carry = "S%d_%d_Carry_%d" % (compress_stage+1, col+1, i)  # this is not the real name which is determined by next col
         comp_carry = "S%d_%d_Carry_%d" % (compress_stage+1, col, i)  # this is not the real name which is determined by next col
         compress_io = [comp_carry, comp_sum] + comp_in # (3,2) counter / full adder
         
         curr_col_result.append(comp_sum)
         cols_carry[col].append(comp_carry)
         compress_array[compress_stage][col].append(compress_io)

         fields.append(3)
      
      
      if n_reminder == 2:
         i += 1
         comp_in = items[i*3:]
         assert(len(comp_in) == 2)
         comp_sum = "S%d_%d_%d" % (compress_stage+1, col, i)
         # comp_carry = "S%d_%d_Carry_%d" % (compress_stage+1, col+1, i)
         comp_carry = "S%d_%d_Carry_%d" % (compress_stage+1, col, i)
         compress_io = [comp_carry, comp_sum] + comp_in # (2,2) counter/half adder

         curr_col_result.append(comp_sum)
         cols_carry[col].append(comp_carry)
         compress_array[compress_stage][col].append(compress_io)

         fields.append(2)

      elif n_reminder == 1:
         i += 1
         comp_in = items[i*3]
         comp_sum = "S%d_%d_%d" % (compress_stage+1, col, i) 
         compress_io = [comp_sum, comp_in]              # just a wire
         
         curr_col_result.append(comp_sum)
         compress_array[compress_stage][col].append(compress_io)
      
         fields.append(1)

      compress_groups.append(fields)

      # calc the previous col's carry out (get its real name)
      if col >= 1:              # ignore column 0, there is no carry
         for j in range(len(cols_carry[col-1])):
            prev_carry = "S%d_%d_%d" % (compress_stage+1, col, i+j+1)
            curr_col_result.append(prev_carry)
            cols_carry[col-1][j] = prev_carry # rename to the real name
            compress_array[compress_stage][col-1][j][0] = prev_carry # rename the previous col's carry out 

      # record the compress result
      for row,v in enumerate(curr_col_result):
         compress_dots[row][col] = (v,False)
      for r in range(row+1, rows_num, 1):
         compress_dots[r][col] = (0,False)

   return compress_groups

def print_compress_sverilog_code():
   stage_num = len(compress_array)
   compress_operations = [[] for i in range(stage_num)]
   create_wires = [[] for i in range(stage_num)]
   FA_cnt = 1
   HA_cnt = 1
   for stage in range(stage_num):
      for col in range(cols_num):
         for item in compress_array[stage][col]:
            # tmp_str = ""
            # for i in item:
            #    if i == '1':        # replace "1" width "1'b1"
            #       tmp_str += "1'b1,"
            #    else:
            #       tmp_str += i+","
            # tmp_str = tmp_str[:-1]
            tmp_str = ','.join(item)
            if len(item) == 5:     # full adder
               compress_operations[stage].append("FA fa%d (%s);" % (FA_cnt, tmp_str))
               create_wires[stage] += item[0:1]
               FA_cnt += 1
            elif len(item) == 4:   # half adder
               compress_operations[stage].append("HA ha%d (%s);" % (HA_cnt, tmp_str))
               create_wires[stage] += item[0:1]
               HA_cnt += 1
            elif len(item) == 2:   # just wire
               # compress_operations[stage].append("assign %s = %s;" % (item[0], tmp_str))
               compress_operations[stage].append("assign %s;" % (tmp_str.replace(",", " = ")))
               create_wires[stage] += [item[0]]
            else:
               assert(False)

   wires_str = ""
   print "// ==================== wires need create ===================="
   for stage in range(stage_num):
      wires_str += "logic " + ','.join(create_wires[stage]) + ";\n"
   # print wires_str

   # compress_str = ""
   compress_str = wires_str
   # compress_str += "logic [%d:0] final_pp0, final_pp1;" % (cols_num-1)
   for stage in range(stage_num):
      compress_str += "// ==================== Stage%d ====================\n" % (stage+1)
      for op in compress_operations[stage]:
         compress_str += op + '\n'

   # the final 2 PPs
   final_pp0 = ""
   final_pp1 = ""
   for col in range(cols_num-1,-1,-1):
      assert(compress_dots[0][col][0]!=0)
      final_pp0 += compress_dots[0][col][0] + ',' # the items of row 0 never has a '0'
      if compress_dots[1][col][0] == 0:
         final_pp1 += "1'b0,"
      else:
         final_pp1 += compress_dots[1][col][0] + ','
   final_pp0 = "{" + final_pp0[:-1] + "}"
   final_pp1 = "{" + final_pp1[:-1] + "}"
   compress_str += "// ==================== Final 2 PPs ====================\n"
   compress_str += "assign final_pp0 = %s;\n" % final_pp0
   compress_str += "assign final_pp1 = %s;\n" % final_pp1

   print compress_str
   print "// -------------------- total %d FAs and %d HAs --------------------" % (FA_cnt-1, HA_cnt-1)
   #   return (FA_cnt-1, HA_cnt-1)


def dot2pos(row, col, pos_base):
   disp_rows_num = rows_num + 1
   point_base = (pos_base[0]+step_x/2, pos_base[1]+step_y/2)
   pos = (point_base[0]+step_x*(cols_num-1-col), point_base[1]+step_y*(disp_rows_num-1-(row+1)))
   return pos
   
# draw oval vertical
def plot_oval(start_pos, end_pos, **kwargs):
   # patch = mpatches.FancyBboxPatch((pos[0]-step_x/2,pos[1]-step_y/2), step_x, 30, fill=False, color="blue", linewidth=3, boxstyle=mpatches.BoxStyle("Round", pad=0))
   if start_pos[0] < end_pos[0]:
      sx = start_pos[0]
      ex = end_pos[0]
   else:
      sx = end_pos[0]
      ex = start_pos[0]

   if start_pos[1] < end_pos[1]:
      sy = start_pos[1]
      ey = end_pos[1]
   else:
      sy = end_pos[1]
      ey = start_pos[1]

   width = ex - sx + 1
   height = ey - sy + 1

   patch = mpatches.FancyBboxPatch((sx, sy), width, height, linewidth=2, boxstyle=mpatches.BoxStyle("Round", pad=3), **kwargs)
   ax.add_patch(patch)


def display_compress_groups(compress_groups, pos_base):
   # draw the select oval
   for col in range(cols_num):
      cnt = 0
      for n in compress_groups[col]:
         if n==2 or n==3:
            spos, epos = dot2pos(cnt, col, pos_base), dot2pos(cnt+n-1, col, pos_base)
            plot_oval(spos, epos, color="red", fill=False)
         cnt += n




curr_pos_index = 0
def gen_next_pos():
   global curr_pos_index
   pos = ( (curr_pos_index/3)*330, 200-(curr_pos_index%3)*100 )
   curr_pos_index += 1
   return pos
   
# # create a new figure and display the dots
# def display_dots_fig(dots, dotcolor="black"):
#    display_dots(dots_1_up, pos_base=(350,0))

####################################################################################################
# the final version for mix multiply 8x8 and 16x16

# pp_8_L[3:0] is the same with pp[3:0]
dots_8_L = gen_pp(mcand_width=8, mplier_width=8, prefix="")
# pp_8_L[4] is not the same with pp[4]
pp_8_L_4 = [("pp_8_L[4][%d]" % i,False) for i in range(8)]
set_pp(dots_8_L, row=4, pp_val=pp_8_L_4, s_val=None)
# add sign compensation when PPs of 8x8 mode(16bits) extend to 16x16 mode(32bits)
for col in range(16, 32, 1):
   set_dots_cell(dots_8_L, 0, col, ("1'b1", True))


# pp_8_H, # not reduce the height, for counteract the "1" below
dots_8_H = gen_pp(mcand_width=8, mplier_width=8, prefix="", pp_index_start=(4,8), pp0_height_reduction=False)


# add dots_8_L and dots_8_H to construct dots_8
dots_8 = copy.deepcopy(dots_8_L)
for col in range(16,31+1,1):
   for row in range(rows_num-1):
      dots_8[row+1][col] = dots_8_H[row][col-16]
# add sign compensation for mix calc(16x16 or two 8x8)
dots_8[rows_num-1][col_middle] = ("C", True)
# reorganize upwards
dots_8_up = reorg_upward(dots_8)

display_dots(dots_8_L, pos_base=gen_next_pos(), style="text")
display_dots(dots_8_H, pos_base=gen_next_pos(), style="text")
display_dots(dots_8, pos_base=gen_next_pos(), style="text")
display_dots(dots_8_up, pos_base=gen_next_pos(), style="text")

# the "1" that could be counteracted
for col in range(25,32,1):
   dots_8_up[0][col] = (0, False)
dots_8_up[2][25] = (0, False)
display_dots(dots_8_up, pos_base=gen_next_pos(), style="text")
dots_8_up = reorg_upward(dots_8_up)
display_dots(dots_8_up, pos_base=gen_next_pos(), style="text")

dots_16 = gen_pp(mcand_width=16, mplier_width=16)
dots_16_up = reorg_upward(dots_16)
display_dots(dots_16_up, pos_base=gen_next_pos(), style="text")

display_dots(dots_8_up, pos_base=gen_next_pos(), style="dot")
display_dots(dots_16_up, pos_base=gen_next_pos(), style="dot", color="red")


# the dots in 8x8 mode is dots_8_up, in 16x16 mode is dots_16_up
pos = gen_next_pos()
display_dots(dots_8_up, pos_base=pos, style="dot")
display_dots(dots_16_up, pos_base=pos, style="dot", color="red", fill=False)

## ========== check the coverage ==========
flag = True
m16_dots_cnt = 0
m8_dots_cnt = 0
for col in range(cols_num):
   for row in range(rows_num):
      if dots_8_up[row][col][0] != 0:
         m8_dots_cnt += 1
      if dots_16_up[row][col][0] != 0:
         m16_dots_cnt += 1
      if (dots_8_up[row][col][0] != 0) and (dots_16_up[row][col][0] == 0):
         print "dots[%d][%d] for 8x8 mode is not covered by 16x16 mode" % (row, col)
         flag = False

if flag == True:
   print "all dots in 8x8 mode is covered by 16x16 mode"

print "dots in 16x16 mode is : %d" % m16_dots_cnt
print "dots in 8x8 mode is : %d" % m8_dots_cnt

# # print the col expression
# print_col_formula(dots_8_up)
# print_col_formula(dots_16_up)
cols_8 = get_col_items(dots_8_up)
cols_16 = get_col_items(dots_16_up)
for items in cols_8:
   items.sort(reverse=True)
for items in cols_16:
   items.sort(reverse=True)

cols_mix = []
for col in range(cols_num):
   col_common = []
   col_16_only = []
   col_8_only = []
   for i in cols_8[col]:
      if i in cols_16[col]:
         cols_16[col].remove(i)
         col_common.append(i)
      else:
         col_8_only.append(i)
   col_16_only = cols_16[col]
   
   # print "col_8[%d] = %s;" % (col, ' + '.join(col_common + col_8_only))
   # print "col_16[%d] = %s;" % (col, ' + '.join(col_common + col_16_only))

   # construct the new column for both 16/8 mode
   assert(len(col_8_only) <= len(col_16_only)) # all dots in 8x8 mode should covered by dots in 16x16 mode
   for i8 in col_8_only:         # each of these signal need a mux ??
      i16 = col_16_only.pop(0)
      col_common.append("(mode8 ? %s : %s)" % (i8,i16))
   for i16 in col_16_only:
      col_common.append("(~mode8 & %s)" % i16)

   cols_mix.append(col_common)

dots_mix = gen_empty_dots_array(rows_num, cols_num)
for col, items in enumerate(cols_mix):
   val_list = [(i,False) for i in items]
   set_dots_col(dots_mix, col, val_list)


display_dots(dots_mix, pos_base=gen_next_pos(), style="text")

# for col,items in enumerate(cols_mix):
#    print "col_mix[%d] = %s;" % (col, ' + '.join(items))
print_col_formula(dots_mix)

# FIXME: print the only and common dots ???


# wallace tree compress
wallace_init(dots_mix)
cnt = 0
while True:
   curr_pos = gen_next_pos()
   display_dots(dots_mix, pos_base=curr_pos, style="text")
   if is_compress_finish():
      break
   
   cnt += 1
   compress_groups = wallace_compress()
   display_compress_groups(compress_groups, pos_base=curr_pos)

print "total need %d stages to compress" % cnt

print_compress_sverilog_code()



print "// ------------------ sum of each column for 16x16 mode --------------------------"
print_col_formula(dots_16_up)
print "// --------------------------------------------------"

print "// ------------------ sum of each column for 8x8 mode --------------------------"
print_col_formula(gen_pp(mcand_width=8, mplier_width=8))
print "// --------------------------------------------------"

show()
exit()



####################################################################################################
# dots_1 = gen_pp([1]*8, 8, 1, prefix="L_")
# dots_2 = gen_pp([1]*8, 8, 1, prefix="H_")

dots_8_L = gen_pp(mcand_width=8, mplier_width=8, prefix="L_")
dots_8_H = gen_pp(mcand_width=8, mplier_width=8, prefix="H_", pp0_height_reduction=False) # not reduce the height, for counteract the "1"

display_dots(dots_8_L, pos_base=gen_next_pos())
display_dots(dots_8_H, pos_base=gen_next_pos())

# # L_8[0] = pp[0][8:0]
# pp_8_L_0 = [("pp[0][%d]" % i,False) for i in range(9)]
# pp_8_L_0 += [("~E[0]",True), ("~E[0]",True), ("E[0]",True)]
# set_pp(dots_8_L, row=0, pp_val=pp_8_L_0, s_val=("S[0]", True))

# set_dots_cell(dots_8_L, row, col, ("K", True))

display_dots(dots_8_L, pos_base=gen_next_pos(), style="text")

pp_8_L_1 = [("pp[1][%d]" % i,False) for i in range(9)]
pp_8_L_1 += [("E[1]",True), ("1'b1",True)]
set_pp(dots_8_L, row=1, pp_val=pp_8_L_1, s_val=("S[1]", True))



display_dots(dots_8_L, pos_base=gen_next_pos(), style="text")
show()
exit()

# compensation for dots_8_L (in the high 16 bits)
for col in range(16,31+1,1):
   dots_8_L[0][col] = ('1', True)

display_dots(dots_8_L, pos_base=gen_next_pos())

# add dots_8_H
for col in range(16,31+1,1):
   for row in range(rows_num-1):
      dots_8_L[row+1][col] = dots_8_H[row][col-16]

dots_8 = dots_8_L
display_dots(dots_8, pos_base=gen_next_pos())

# add sign compensation for mix calc(16x16 or two 8x8)
dots_8[rows_num-1][col_middle] = ("C", True)
display_dots(dots_8, pos_base=gen_next_pos())

dots_8_up = reorg_upward(dots_8)
display_dots(dots_8_up, pos_base=gen_next_pos())

# the "1" that could be counteracted
for col in range(25,32,1):
   dots_8_up[0][col] = (0, False)
dots_8_up[2][25] = (0, False)
display_dots(dots_8_up, pos_base=gen_next_pos())

dots_8_up = reorg_upward(dots_8_up)
display_dots(dots_8_up, pos_base=gen_next_pos())

dots_16 = gen_pp(mcand_width=16, mplier_width=16, prefix="M16_")
dots_16_up = reorg_upward(dots_16)
display_dots(dots_16_up, pos_base=gen_next_pos())

# show all dots as a filled circle, never show text
for row in range(rows_num):
   for col in range(cols_num):
      # for dots_16_up
      if dots_16_up[row][col][0] != 0:
         dots_16_up[row][col] = (dots_16_up[row][col][0], False)
      # for dots_8_up
      if dots_8_up[row][col][0] != 0:
         dots_8_up[row][col] = (dots_8_up[row][col][0], False)
         
pos = gen_next_pos()
display_dots(dots_8_up, pos_base=pos)
display_dots(dots_16_up, pos_base=pos, color="red", fill=False)


# print_col_formula(dots_8_L)
# print_col_formula(dots_8_H)

show()
exit()


# pos = gen_next_pos()

display_dots(dots_8, pos_base=gen_next_pos())
display_dots(dots_8_H, pos_base=gen_next_pos(), style="dot")
display_dots(dots_8_L, pos_base=gen_next_pos(), style="text")
display_dots(dots_8, pos_base=gen_next_pos(), fill=False)
display_dots(dots_16, pos_base=gen_next_pos(), dotstyle="halfcircle", color="red")


pp_8_L_0 = [("pp[0][%d]" % i,False) for i in range(9)]
pp_8_L_0 += [("~E[0]",True), ("~E[0]",True), ("E[0]",True)]
set_pp(dots_8_L, row=0, pp_val=pp_8_L_0, s_val=("S[0]", True))

set_dots_cell(dots_8_L, row=1, col=14, val=("TEST", True))

display_dots(dots_8_L, pos_base=gen_next_pos(), style="text")

# for col in range(cols_num):
#    for row in range(rows_num):

print_col_formula(dots_1_up)
print_col_formula(dots_16_up)


savefig("tmp.png", dpi=72)
show()
exit()
